import os
import pandas as pd

REPLAY_ATTACK_SAVE_FOLDER = '../data/replay_attack/'


def generate_csv(root_dir, mode):
    if 'train' in mode:
        videos_path = os.path.join(root_dir, 'replayattack-train_crops')
        # videos_path = os.path.join(root_dir, 'Train_files_MSR')
    else:
        videos_path = os.path.join(root_dir, 'replayattack-test_crops')
        # videos_path = os.path.join(root_dir, 'Test_files_MSR')
    fold_data = []
    for video_name in os.listdir(videos_path):
        video_name_split = video_name.split('_')
        video_path = os.path.join(videos_path, video_name)
        for image in os.listdir(video_path):
            if len(video_name_split) == 7:
                label = 0
                fold_data.append([video_name, image, label])
            elif len(video_name_split) == 8:
                label = 1
                fold_data.append([video_name, image, label])

    columns = ["video", "file", "label"]
    # pd.DataFrame(fold_data, columns=columns).to_csv(f'{REPLAY_ATTACK_SAVE_FOLDER}data_{mode}.csv', index=False)
    # 1_frame
    pd.DataFrame(fold_data, columns=columns).to_csv(f'{REPLAY_ATTACK_SAVE_FOLDER}data_{mode}_1_frame.csv', index=False)
    # pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}_MSR.csv'.format(mode), index=False)


def main():
    os.makedirs(REPLAY_ATTACK_SAVE_FOLDER, exist_ok=True)
    # root_dir = '/data/heyan/Datasets/Face/oulu_256'  # RGB images
    # root_dir = '/home/shaohua/data2/Datasets/Oulu_NPU_256'  # MSR images
    # root_dir = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Replay_Attack/Replayattack/replayattack'
    # 1_frame
    root_dir = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Replay_Attack/Replayattack_1_frame/replayattack'

    generate_csv(root_dir, 'train')
    # generate_csv(root_dir, 'val')
    generate_csv(root_dir, 'test')

    # generate_csv(root_dir, 'train_real')
    # generate_csv(root_dir, 'train_fake')


if __name__ == '__main__':
    main()
